import pandas as pd
import numpy as np
import matplotlib.pylab as plt
plt.rcParams.update({'font.size': 22})
plt.rcParams.update({'figure.figsize': (14, 6)})
plt.rcParams['font.family'] = 'serif'
plt.rcParams['font.sans-serif'] = ['Verdana']
plt.rcParams['xtick.major.pad']='8'
import seaborn as sns

def make_fig_1_plot(res, B):
    #=== Make the plot
    fig, (ax, ax2) = plt.subplots(2, 1, figsize=[12, 14])

    
    l =res[['Bias_X', 'Bias_Z']].unstack().reset_index()
    l['level_0'] = l['level_0'].replace(['Bias_X', 'Bias_Z'], [r'$\rho = \alpha$', r'$\rho = \gamma$'])
    sns.boxplot(x = 'level_0', y = 0, data = l, showfliers = False, 
                showmeans=False, palette = ['navy', 'r'],
                meanprops={"marker":"o",
                           "markerfacecolor":"red", 
                           "markeredgecolor":"black",
                          "markersize":"10"}, ax = ax)
    ax.set_xlabel('', fontsize = 36)
    ax.set_ylabel(r'Bias, $\frac{\hat{\rho} - \rho}{\rho}$', fontsize = 32)
    ax.axhline(0, c = 'k', lw = 2)
    
    
    B.columns = [r'$\psi$', 'Opposite sign', 'Not significant']
    m = B.set_index(r'$\psi$').unstack().reset_index()
    m[r'$\psi$'] = m[r'$\psi$'].round(2)
    m = m.rename(columns = {'level_0': ''})
    K = sns.boxplot(x = r'$\psi$', y = 0, hue = '', data = m, showfliers = False,palette = ['#DE970B', '#808080'])
    plt.ylabel('Proportion of coefficients $\gamma$', fontsize = 26)
    for patch in K.patches:
        r, g, b, a = patch.get_facecolor()
        patch.set_facecolor((r, g, b, .8))
    plt.savefig('ModelA.pdf', bbox_inches = 'tight')
    
def make_fig_2_plot(res, dx, z):
    #=== Make the plot
    fig, (ax, ax2) = plt.subplots(2, 1, figsize=[12, 14])
    ax = plt.subplot(211)

    l =res[['Bias_X', 'Bias_XB']].unstack().reset_index()
    l['level_0'] = l['level_0'].replace(['Bias_X', 'Bias_XB'], ['Excluding $C$', 'Including $C$'])
    sns.boxplot(x = 'level_0', y = 0, data = l, showfliers = False, showmeans=False, palette = ['teal', 'darkorange'],
                meanprops={"marker":"o",
                           "markerfacecolor":"red", 
                           "markeredgecolor":"black",
                          "markersize":"10"}, ax = ax)
    plt.xlabel('', fontsize = 36)
    plt.xticks(fontsize=24)
    plt.ylabel(r'Bias, $\frac{\hat{\alpha} - \alpha}{\alpha}$', fontsize = 36)
    plt.axhline(0, c = 'k', lw = 2)
    ax.set_ylim([-0.6,0.2])
    left, bottom, width, height = [0.25, 0.6, 0.18, 0.12]
    axA = fig.add_axes([left, bottom, width, height])
    res[['R2', 'R2C']].plot.hist(ax = axA, color = ['teal', 'darkorange'], bins = 100, legend = False)
    axA.set_xlabel(r'R$^2$', loc = 'center', labelpad=0)
    axA.set_ylabel('')
    #====
    dx_ = pd.DataFrame(dx, columns = [r'$\eta$', 'bias', 'bias_std', 'negative']).set_index(r'$\eta$')
    sns.boxplot(x = r'$\eta$', y = 'Bias_XB', data = z, palette = 'RdBu_r', showfliers = False, ax = ax2)
    ax2.set_ylabel(r'Bias, $\frac{\hat{\alpha} - \alpha}{\alpha}$', fontsize = 36)
    plt.savefig('ModelB.pdf', bbox_inches = 'tight')

def make_fig_3_plot(res):
    plt.figure(figsize = (12,16))
    ax = plt.subplot(211)
    est = res[res.columns[:3]]
    est.columns = [r'$\alpha: P\to L$', r'$\psi: S\to L$ (partial)', r'$\sigma: S\to L$ (total)']
    
    est = est.unstack().reset_index().drop(columns = ['level_1'])
    sns.barplot(x='level_0', y = 0, data = est, palette=['tab:cyan', 'darkorange', 'forestgreen'], ax= ax, ci=99, capsize = .2)
    plt.xlabel('')
    plt.ylabel('Regression coefficient', fontsize = 26)
    plt.savefig('ModelC.pdf', bbox_inches = 'tight')
